package hku.cs.lunuo.instrument;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import soot.Body;
import soot.BodyTransformer;
import soot.Local;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Unit;
import soot.jimple.IdentityStmt;
import soot.jimple.IfStmt;
import soot.jimple.IntConstant;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.Jimple;
import soot.jimple.LookupSwitchStmt;
import soot.jimple.ReturnStmt;
import soot.jimple.ReturnVoidStmt;
import soot.jimple.StringConstant;
import soot.jimple.TableSwitchStmt;
import soot.toolkits.graph.Block;
import soot.toolkits.graph.BlockGraph;
import soot.toolkits.graph.ExceptionalBlockGraph;

public class PredAndStmtInstrument2 extends BodyTransformer {

	private static PredAndStmtInstrument2 instance = new PredAndStmtInstrument2();

	private PredAndStmtInstrument2() {
	}

	public static PredAndStmtInstrument2 v() {
		return instance;
	}

	SootClass counterClass;
	SootMethod increaseES, flushES, increaseCounter, add, recordLast,
			confirmLast;

	@Override
	protected void internalTransform(Body b, String arg1, Map arg2) {
		counterClass = Scene.v()
				.getSootClass("hku.cs.lunuo.instrument.Counter");
		increaseES = counterClass
				.getMethod("void increaseES(java.lang.String,java.lang.String,int)");
		increaseCounter = counterClass
				.getMethod("void increase(java.lang.String,java.lang.String,int)");
		flushES = counterClass
				.getMethod("void flushES(java.lang.String,java.lang.String,java.lang.String,int,int)");
		add = counterClass
				.getMethod("void add(java.lang.String,java.lang.String,java.lang.String)");
		recordLast = counterClass
				.getMethod("void recordLast(java.lang.String,java.lang.String,int)");
		confirmLast = counterClass.getMethod("void confirmLast()");

		BlockGraph bGraph = new ExceptionalBlockGraph(b);

		// System.out.println("wo shi xin de ");

		Iterator<Block> blockIt = bGraph.getBlocks().iterator();
		Local counterLocal = Jimple.v().newLocal("counterRef",
				counterClass.getType());
		b.getLocals().add(counterLocal);
		String methodName = b.getMethod().toString();
		String className = methodName.substring(1, methodName.indexOf(":"));
		methodName = methodName.substring(methodName.indexOf(":") + 2,
				methodName.length() - 1);

		List<Block> blockLs = bGraph.getBlocks();
		Set<Integer> hasRecord = new HashSet<Integer>();
		boolean globalNoBranch = true;
//		boolean xS = false;
//		boolean yS = false;
//		String fix = "-1";
//		for (int xi = 0; xi < blockLs.size(); xi++) {
//			Block block = blockLs.get(xi);
//			Unit ut = block.getTail();
//			
//			if(block.getPreds().size()>1)
//				xS = true;
//			
//			if(ut.branches())
//				yS = true; 
//			
//			if(xS || yS) {
//				if(xS&&!yS) {
//					List argxs = new ArrayList();
//					argxs.add(StringConstant.v(className));
//					argxs.add(StringConstant.v(methodName));
//					argxs.add(StringConstant.v(fix));
//					InvokeExpr incExpr2 = Jimple.v().newStaticInvokeExpr(
//							add.makeRef(), argxs);
//					InvokeStmt incStmt2 = Jimple.v().newInvokeStmt(incExpr2);
//					Unit utu = blockLs.get(0).getTail();
//					if(!(utu instanceof IdentityStmt))
//						blockLs.get(0).insertBefore(incStmt2, utu);
//					else
//						blockLs.get(0).insertAfter(incStmt2, utu);
//				}
//				break;
//			}
//			
//			fix+=">"+xi;
//		}
		
		for (int bi = 0; bi < blockLs.size(); bi++) {
			Block block = blockLs.get(bi);
			Unit ut = block.getTail();
			boolean shouldCatch = false;
			int bj = bi;
			if (ut instanceof IfStmt) {
				while (bj < blockLs.size()
						&& blockLs.get(bj + 1).getTail() instanceof IfStmt) {
					List<Block> preBls = blockLs.get(bj + 1).getPreds();
					boolean isNotCon = false;
					for (Block bses : preBls)
						if (bses.getIndexInMethod() < bi
								|| bses.getIndexInMethod() > bj) {
							isNotCon = true;
							break;
						}

					if (!isNotCon)
						bj++;
					else
						break;

				}
				shouldCatch = true;
			} else if (ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt) {
				shouldCatch = true;
			}

			if (shouldCatch && block.getPreds().size() >= 2) {
//				if(!(block.getHead() instanceof IdentityStmt)) {
				InvokeExpr incExpr = Jimple.v().newStaticInvokeExpr(
						confirmLast.makeRef());
				InvokeStmt incStmt = Jimple.v().newInvokeStmt(incExpr);
				Unit uh = block.getHead();
				block.insertBefore(incStmt, uh);
				uh.redirectJumpsToThisTo(incStmt);
//				}

				for (Block preb : block.getPreds()) {
					int prebi = preb.getIndexInMethod();
					if (!hasRecord.contains(prebi)) {
						hasRecord.add(prebi);
						List args = new ArrayList();
						args.add(StringConstant.v(className));
						args.add(StringConstant.v(methodName));
						args.add(IntConstant.v(prebi));
						InvokeExpr incExpr2 = Jimple.v().newStaticInvokeExpr(
								recordLast.makeRef(), args);
						InvokeStmt incStmt2 = Jimple.v()
								.newInvokeStmt(incExpr2);

						Unit iut = preb.getTail();

						if (preb.getTail().branches()) {
							preb.insertBefore(incStmt2, iut);
							iut.redirectJumpsToThisTo(incStmt2);
						} else {
							preb.insertAfter(incStmt2, iut);
						}// actually the following one is also ok;

						// if (preb.getTail() instanceof
						// IdentityStmt||!preb.getTail().branches()) {
						// preb.insertAfter(incStmt2, iut);
						// } else {
						// preb.insertBefore(incStmt2, iut);
						// ut.redirectJumpsToThisTo(incStmt2);
						// }
						// it is right now to use like this by adding
						// "!preb.getTail().branches()"
						// for instance, there is one special condition
						// like method in "AbstractTokenizer" namely
						// "completeWhitespace"
						// Block 0:[preds:] [succs: 1]
						// 944: l0 := @this:
						// de.susebox.java.util.AbstractTokenizer
						// 944: l1 := @parameter0: de.susebox.java.util.Token
						// 944: $i0 =
						// l0.<de.susebox.java.util.AbstractTokenizer: int
						// _currentReadPos>
						// 944: l2 = $i0 + 1
						// 945: $i0 =
						// l0.<de.susebox.java.util.AbstractTokenizer: int
						// _currentWritePos>
						// 945: l3 = $i0 - l2
						// 946: $i1 =
						// l0.<de.susebox.java.util.AbstractTokenizer: int
						// _rangeStart>
						// 946: $i1 = $i1 + l2
						// 946: l4 = virtualinvoke
						// l0.<de.susebox.java.util.AbstractTokenizer: int
						// readWhitespaces(int,int)>($i1, l3)
						//
						// Block 1:[preds: 0 4] [succs: 2 5]
						// 948: if l4 != l3 goto virtualinvoke
						// l1.<de.susebox.java.util.Token: void setType(int)>(6)
						//
						// Block 2:[preds: 1] [succs: 3 4]
						// 949: $i0 = virtualinvoke
						// l0.<de.susebox.java.util.AbstractTokenizer: int
						// readMoreData()>()
						// 949:
						// here is a "while grammer" in source code then we can
						// see the tail unit of block 0 is
						// not a branch statement, then the recordLast()
						// statement will lost because next statement is not
						// confirm() statemnet but a function call, then .....

					}
				}
			}
			bi = bj;
		}

		for (int i = 0; i < blockLs.size(); i++) {
			Unit ut = blockLs.get(i).getTail();
			if (ut instanceof IfStmt || ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt) {
				globalNoBranch = false;
				String startIs = "";
				List<Block> preBs = blockLs.get(i).getPreds();

				if (preBs.size() < 2)
					startIs = String.valueOf(i);
				else {
					for (Block preb : preBs) {
						startIs += "_"
								+ String.valueOf(preb.getIndexInMethod());
					}
					startIs = startIs.substring(1, startIs.length());
				}

				if (ut instanceof IfStmt) {
					int j = i;
					while (j < blockLs.size()
							&& blockLs.get(j + 1).getTail() instanceof IfStmt) {
						List<Block> preBls = blockLs.get(j + 1).getPreds();
						boolean sb = false;
						for (Block bsss : preBls)
							if (bsss.getIndexInMethod() < i
									|| bsss.getIndexInMethod() > j) {
								sb = true;
								break;
							}

						if (!sb)
							j++;
						else
							break;

					}
					int k = i;
					while (k <= j) {
						if (k == j) {
							Unit uh = blockLs.get(k + 1).getHead();

							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(k));
							args.add(IntConstant.v(-(k + 1)));// here make
																// negative
																// value means
																// false branch
							InvokeExpr incExpr1 = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt1 = Jimple.v().newInvokeStmt(
									incExpr1);
							blockLs.get(k + 1).insertBefore(incStmt1, uh);
						}

						List<Block> sucBlks = blockLs.get(k).getSuccs();
						List<Integer> sucRightBlock = new ArrayList<Integer>();
//						if (className.contains("AbstractTokenizer")
//								&& methodName.contains("nextToken") && k == 12)
//							System.out.println();

						for (Block sucb : sucBlks) {

							if (!(sucb.getHead() instanceof IdentityStmt)) {
								int targetBlkIdx = sucb.getIndexInMethod();
								if (targetBlkIdx != k + 1) {
									sucRightBlock.add(targetBlkIdx);
								}
							}
							// never ifstmt goto first block which contain
							// IdentityStmt
							// sometime the code will like this
							/*
							 * Block 0:[preds:] [succs: 2] ~~~~~~~~~~: l0 :=
							 * @this: de.susebox.java.util.AbstractTokenizer
							 * ~~~~~~~~~~: l1 := @parameter0:
							 * de.susebox.java.util.AbstractTokenizer 832: goto
							 * [?= staticinvoke
							 * <hku.cs.lunuo.instrument.Counter: void
							 * increase(java.lang.String,java.lang.String,int)>(
							 * "de.susebox.java.util.AbstractTokenizer",
							 * "de.susebox.java.util.AbstractTokenizer getBaseTokenizer(de.susebox.java.util.AbstractTokenizer)"
							 * , 2)]
							 * 
							 * Block 1:[preds: 2] [succs: 2] 833: l1 =
							 * l1.<de.susebox.java.util.AbstractTokenizer:
							 * de.susebox.java.util.AbstractTokenizer
							 * _prevTokenizer>
							 * 
							 * Block 2:[preds: 0 1] [succs: 3 1] 832: l2 =
							 * l1.<de.susebox.java.util.AbstractTokenizer:
							 * de.susebox.java.util.AbstractTokenizer
							 * _prevTokenizer> 832: if l2 != null goto
							 * staticinvoke <hku.cs.lunuo.instrument.Counter:
							 * void
							 * increase(java.lang.String,java.lang.String,int
							 * )>("de.susebox.java.util.AbstractTokenizer",
							 * "de.susebox.java.util.AbstractTokenizer getBaseTokenizer(de.susebox.java.util.AbstractTokenizer)"
							 * , 1)
							 * 
							 * Block 3:[preds: 2] [succs:] 835: return l1
							 */
						}

						int targetBlkIdx;

						if (sucRightBlock.size() > 1) {
							System.out.println("current block " + k);
							System.out.println("suc Block " + sucRightBlock);
							System.out.println(className + " " + methodName);
							System.out.println(sucBlks);
							throw new IllegalStateException(
									"must not have more than two suc normal block");
						} else if (sucRightBlock.size() == 0)
							targetBlkIdx = k + 1;// deal empty if block or while
						// block, may be this will
						// has no true branch
						else
							targetBlkIdx = sucRightBlock.get(0);

						// targetBlkIdx <= i because while structure will loop
						// in one block
						// Block 1:
						// [preds: 0 1] [succs: 2 1]
						// 159:
						// 160: if l4 != -1 goto <159>
						if (targetBlkIdx > j || targetBlkIdx <= i) {
							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(k));
							args.add(IntConstant.v(targetBlkIdx));// here use
																	// object
																	// value
																	// means
																	// true
																	// branch

							Block instruBlk = blockLs.get(targetBlkIdx);
							InvokeExpr incExpr = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt = Jimple.v().newInvokeStmt(
									incExpr);
							Unit uh = instruBlk.getHead();

//							 if (!(uh instanceof IdentityStmt)) {
//							 instruBlk.insertBefore(incStmt, uh);
//							 uh.redirectJumpsToThisTo(incStmt);
//							 } else {
//							 Unit tmp = uh;
//							 while (tmp != null && tmp instanceof
//							 IdentityStmt) {
//							 tmp = instruBlk.getSuccOf(tmp);
//							 }
//							 if (tmp != null)
//							 instruBlk.insertBefore(incStmt, tmp);
//							 else
//							 instruBlk.insertAfter(incStmt,
//							 instruBlk.getTail());
//							 }

							instruBlk.insertBefore(incStmt, uh);
							uh.redirectJumpsToThisTo(incStmt);

						}
						k++;
					}
					i = j;
				} else if (ut instanceof TableSwitchStmt
						|| ut instanceof LookupSwitchStmt) {
					Block block = blockLs.get(i);
					List<Block> sucBs = block.getSuccs();

					for (Block sucb : sucBs) {
						Block instruBlk = sucb;
						Unit uh = instruBlk.getHead();
						if (!(uh instanceof IdentityStmt)) {
							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(i));
							args.add(IntConstant.v(sucb.getIndexInMethod()));
							// here we don't differ the default case and other
							// case

							InvokeExpr incExpr = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt = Jimple.v().newInvokeStmt(
									incExpr);

							instruBlk.insertBefore(incStmt, uh);
							uh.redirectJumpsToThisTo(incStmt);
						}
					}
				}
			}
		}
		
		Block bxx = blockLs.get(0);
		Unit uxt = bxx.getTail();
		Unit uxh = bxx.getHead();
		List argxs = new ArrayList();
		argxs.add(StringConstant.v(className));
		argxs.add(StringConstant.v(methodName));
		argxs.add(StringConstant.v("-1>0"));
		InvokeExpr incExpr3 = Jimple.v().newStaticInvokeExpr(
				add.makeRef(), argxs);
		InvokeStmt incStmt3 = Jimple.v().newInvokeStmt(incExpr3);
		
//		if(ux instanceof IdentityStmt) {
//			bxx.insertAfter(incStmt3, ux);
//		}else {
//			ux.redirectJumpsToThisTo(incStmt3);
//			bxx.insertBefore(incStmt3, ux);
//		}
		
		
		Unit tmxp = uxh;
		while (tmxp != null && tmxp instanceof IdentityStmt) {
			tmxp = bxx.getSuccOf(tmxp);
		}
		if (tmxp != null)
			bxx.insertBefore(incStmt3, tmxp);
		else
			bxx.insertAfter(incStmt3, uxt);
		

		while (blockIt.hasNext()) {
			Block block = (Block) blockIt.next();
			int blkIdx = block.getIndexInMethod();

			List args = new ArrayList();
			args.add(StringConstant.v(className));
			args.add(StringConstant.v(methodName));
			args.add(IntConstant.v(blkIdx));

			InvokeExpr incExpr1 = Jimple.v().newStaticInvokeExpr(
					increaseCounter.makeRef(), args);
			InvokeStmt incStmt1 = Jimple.v().newInvokeStmt(incExpr1);
			// must insirt in head but after identityStmt becuase if you
			// insirt before tail, when run this block, the stmt before the
			// counter stmt throw exception, then this counter will not be
			// recorded

			Unit ut = block.getTail();
			Unit uh = block.getHead();

			if (!(uh instanceof IdentityStmt)) {
				block.insertBefore(incStmt1, uh);
				ut.redirectJumpsToThisTo(incStmt1);
			} else {
				Unit tmp = uh;
				while (tmp != null && tmp instanceof IdentityStmt) {
					tmp = block.getSuccOf(tmp);
				}
				if (tmp != null)
					block.insertBefore(incStmt1, tmp);
				else
					block.insertAfter(incStmt1, ut);
			}
			
//			if(!(ut instanceof IdentityStmt)) {
//				block.insertBefore(incStmt1, ut);
//				ut.redirectJumpsToThisTo(incStmt1);
//			}else
//				block.insertAfter(incStmt1, ut);
			

			if (ut instanceof IfStmt || ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt) {
				// || ut instanceof ReturnStmt || ut instanceof ReturnVoidStmt)
				// {
				InvokeExpr incExpr2 = Jimple.v().newStaticInvokeExpr(
						increaseES.makeRef(), args);
				InvokeStmt incStmt2 = Jimple.v().newInvokeStmt(incExpr2);
				block.insertBefore(incStmt2, ut);
			} 
//			else if (globalNoBranch
//					&& (ut instanceof ReturnStmt || ut instanceof ReturnVoidStmt)) {
//				List argxs = new ArrayList();
//				argxs.add(StringConstant.v(className));
//				argxs.add(StringConstant.v(methodName));
//				argxs.add(StringConstant.v("-2>" + String.valueOf(blkIdx)));
//				// use -2 becuase maybe there is no ifstmt in method, but switch
//				// is also a branch,
//				// this will make return block has more than 1 pred block, so -2
//				// mean breath first search
//				InvokeExpr incExpr2 = Jimple.v().newStaticInvokeExpr(
//						add.makeRef(), argxs);
//				InvokeStmt incStmt2 = Jimple.v().newInvokeStmt(incExpr2);
//				block.insertAfter(incStmt2, incStmt1);
//			}
		}
	}
	
	private void insertUnitToHead(Unit insertUnit, Block insrtBlock) {
		Unit uh = insrtBlock.getHead();
		
		if (!(uh instanceof IdentityStmt)) {
			insrtBlock.insertBefore(insertUnit, uh);
			uh.redirectJumpsToThisTo(insertUnit);
		} else {
			Unit tmp = uh;
			while (tmp != null && tmp instanceof IdentityStmt) {
				tmp = insrtBlock.getSuccOf(tmp);
			}
			if (tmp != null)
				insrtBlock.insertBefore(insertUnit, tmp);
			else
				insrtBlock.insertAfter(insertUnit, insrtBlock.getTail());
		}
	}
	
}